# -*- coding: utf-8 -*-
"""
Created on Fri May 24 16:45:17 2024

@author: Yu
"""

import sys
import os
import configparser
import socket
import time
from PyQt5.QtWidgets import (
    QApplication, QMainWindow, QTableView, QAbstractItemView, QPushButton,
    QVBoxLayout, QWidget, QHBoxLayout, QLabel, QCheckBox, QMessageBox, QLineEdit,
    QDialog, QFormLayout, QDialogButtonBox, QTabWidget
)
from PyQt5.QtCore import QAbstractTableModel, Qt, QThread, pyqtSignal
from scapy.all import ARP, Ether, srp
import netifaces
from wakeonlan import send_magic_packet
from fabric import Connection
from cryptography.fernet import Fernet
import logging
from threading import Thread


# Setup logging
logging.basicConfig(filename='network_tool.log', level=logging.DEBUG,
                    format='%(asctime)s:%(levelname)s:%(message)s')

def load_key():
    key_path = "secret.key"
    if not os.path.exists(key_path):
        key = Fernet.generate_key()
        with open(key_path, 'wb') as key_file:
            key_file.write(key)
    else:
        with open(key_path, 'rb') as key_file:
            key = key_file.read()
    return key

def encrypt_password(password, key):
    fernet = Fernet(key)
    return fernet.encrypt(password.encode()).decode()

def decrypt_password(encrypted_password, key):
    fernet = Fernet(key)
    return fernet.decrypt(encrypted_password.encode()).decode()

class ScanTableModel(QAbstractTableModel):
    checkboxStateChanged = pyqtSignal()

    def __init__(self, headers):
        super().__init__()
        self.headers = headers
        self.data = []

    def rowCount(self, parent=None):
        return len(self.data)

    def columnCount(self, parent=None):
        return len(self.headers)

    def data(self, index, role):
        if role == Qt.DisplayRole:
            if index.column() == 1:
                return self.data[index.row()][1]
            elif index.column() == 2:
                return self.data[index.row()][2]
            elif index.column() == 3:
                return self.data[index.row()][3]
        elif role == Qt.CheckStateRole and index.column() == 0:
            return self.data[index.row()][0].checkState()
        return None

    def headerData(self, section, orientation, role):
        if role == Qt.DisplayRole:
            if orientation == Qt.Horizontal:
                return self.headers[section]
            elif orientation == Qt.Vertical:
                return str(section + 1)
        return None

    def flags(self, index):
        if index.column() == 0:
            return Qt.ItemIsUserCheckable | Qt.ItemIsEnabled
        return super().flags(index)

    def setData(self, index, value, role):
        if role == Qt.CheckStateRole and index.column() == 0:
            self.data[index.row()][0].setCheckState(value)
            self.dataChanged.emit(index, index)
            self.checkboxStateChanged.emit()
            return True
        return False

class ScanThread(QThread):
    finished = pyqtSignal(list)

    def __init__(self):
        super().__init__()

    def run(self):
        local_ip = socket.gethostbyname(socket.gethostname())
        gateway_ip = self.get_default_gateway()
        ip_range = local_ip + '/24'
        clients = self.scan_network(ip_range, local_ip, gateway_ip, retries=5)
        clients.sort(key=lambda x: x['ip'])
        self.finished.emit(clients)

    def get_default_gateway(self):
        gws = netifaces.gateways()
        return gws['default'][netifaces.AF_INET][0] if 'default' in gws and netifaces.AF_INET in gws['default'] else None

    def scan_network(self, ip_range, local_ip, gateway_ip, retries=3):
        clients = []
        for _ in range(retries):
            arp_request = ARP(pdst=ip_range)
            broadcast = Ether(dst="ff:ff:ff:ff:ff:ff")
            arp_request_broadcast = broadcast / arp_request
            answered_list = srp(arp_request_broadcast, timeout=2, verbose=False)[0]

            for sent, received in answered_list:
                hostname = "N/A"
                if received.psrc == local_ip:
                    hostname = "本机"
                elif received.psrc == gateway_ip:
                    hostname = "网关"
                else:
                    try:
                        hostname = socket.getfqdn(received.psrc)
                    except socket.error:
                        hostname = "N/A"
                clients.append({'ip': received.psrc, 'mac': received.hwsrc, 'hostname': hostname})

            if clients:
                break
            else:
                time.sleep(1)

        return clients

class SettingsDialog(QDialog):
    def __init__(self, parent=None):
        super().__init__(parent)
        self.setWindowTitle("设置-Settings")
        self.setGeometry(300, 300, 300, 200)
        self.key = load_key()
        self.initUI()

    def initUI(self):
        self.tab_widget = QTabWidget()

        # SSH Settings tab
        self.ssh_tab = QWidget()
        self.ssh_layout = QFormLayout()
        self.ssh_username_line_edit = QLineEdit()
        self.ssh_password_line_edit = QLineEdit()
        self.ssh_password_line_edit.setEchoMode(QLineEdit.Password)  # 设置为密码模式
        self.ssh_layout.addRow("SSH 用户名:", self.ssh_username_line_edit)
        self.ssh_layout.addRow("SSH 密码:", self.ssh_password_line_edit)
        self.ssh_tab.setLayout(self.ssh_layout)

        # About tab
        self.about_tab = QWidget()
        self.about_layout = QVBoxLayout()
        self.about_text = QLabel("菁云网络开关机助手 V1.0.0。\n\n有任何问题请随时联系我：www.tyab.cn")
        self.about_layout.addWidget(self.about_text)
        self.about_tab.setLayout(self.about_layout)

        # 加标签到标签窗口
        self.tab_widget.addTab(self.ssh_tab, "SSH 设置")
        self.tab_widget.addTab(self.about_tab, "关于")

        # 按钮窗口
        self.button_box = QDialogButtonBox(QDialogButtonBox.Save | QDialogButtonBox.Cancel)
        self.button_box.accepted.connect(self.save_settings)
        self.button_box.rejected.connect(self.reject)

        # 总布局
        self.layout = QVBoxLayout()
        self.layout.addWidget(self.tab_widget)
        self.layout.addWidget(self.button_box)
        self.setLayout(self.layout)
        self.load_settings()

    def load_settings(self):
        config = configparser.ConfigParser()
        if os.path.exists("settings.ini"):
            config.read("settings.ini")
            self.ssh_username_line_edit.setText(config.get("SSH", "username", fallback=""))
            encrypted_password = config.get("SSH", "password", fallback="")
            if encrypted_password:
                decrypted_password = decrypt_password(encrypted_password, self.key)
                self.ssh_password_line_edit.setText(decrypted_password)

    def save_settings(self):
        encrypted_password = encrypt_password(self.ssh_password_line_edit.text(), self.key)
        config = configparser.ConfigParser()
        config["SSH"] = {
            "username": self.ssh_username_line_edit.text(),
            "password": encrypted_password
        }
        with open("settings.ini", "w") as configfile:
            config.write(configfile)
        QMessageBox.information(self, "保存成功", "设置已成功保存。")

class MainWindow(QMainWindow):
    def __init__(self):
        super().__init__()
        self.setWindowTitle("菁云 网络批量开关机助手 V1.0.0")
        self.setGeometry(100, 100, 500, 400)
        self.initUI()

    def initUI(self):
        self.table_model = ScanTableModel(["Select", "IP Address", "MAC Address", "Hostname"])
        self.table_view = QTableView()
        self.table_view.setModel(self.table_model)
        self.table_view.setSelectionBehavior(QAbstractItemView.SelectRows)

        self.table_view.setColumnWidth(0, 45)
        self.table_view.setColumnWidth(1, 100)
        self.table_view.setColumnWidth(2, 150)
        self.table_view.setColumnWidth(3, 150)

        self.scan_button = QPushButton("扫描网络")
        self.scan_button.clicked.connect(self.start_scan)

        self.select_all_button = QPushButton("全选")
        self.select_all_button.clicked.connect(self.select_all)

        self.deselect_all_button = QPushButton("取消全选")
        self.deselect_all_button.clicked.connect(self.deselect_all)

        self.wake_on_lan_button = QPushButton("开机")
        self.wake_on_lan_button.clicked.connect(self.wake_on_lan)

        self.shutdown_button = QPushButton("关机")
        self.shutdown_button.clicked.connect(self.shutdown)

        self.restart_button = QPushButton("重启")
        self.restart_button.clicked.connect(self.restart)

        self.settings_button = QPushButton("设置")
        self.settings_button.clicked.connect(self.open_settings)

        self.save_button = QPushButton("保存")
        self.save_button.clicked.connect(self.save_data)

        button_layout1 = QHBoxLayout()
        button_layout1.addWidget(self.scan_button)
        button_layout1.addWidget(self.select_all_button)
        button_layout1.addWidget(self.deselect_all_button)
        button_layout1.addWidget(self.save_button)
        button_layout1.addStretch()

        button_layout2 = QHBoxLayout()
        button_layout2.addWidget(self.wake_on_lan_button)
        button_layout2.addWidget(self.shutdown_button)
        button_layout2.addWidget(self.restart_button)
        button_layout2.addWidget(self.settings_button)
        button_layout2.addStretch()

        label_layout = QHBoxLayout()
        self.ip_count_label = QLabel("IP 数量: 0")
        self.selected_count_label = QLabel("选择的数量: 0")
        self.status_label = QLabel("状态: Ready")
        label_layout.addWidget(self.ip_count_label)
        label_layout.addWidget(self.selected_count_label)
        label_layout.addWidget(self.status_label)
        label_layout.addStretch()

        layout = QVBoxLayout()
        layout.addLayout(button_layout1)
        layout.addLayout(button_layout2)
        layout.addWidget(self.table_view)
        layout.addLayout(label_layout)

        central_widget = QWidget()
        central_widget.setLayout(layout)
        self.setCentralWidget(central_widget)

        self.load_selected_ips()
        self.load_settings()

        self.scan_thread = ScanThread()
        self.scan_thread.finished.connect(self.display_results)
        self.table_model.checkboxStateChanged.connect(self.update_selected_count_label)

    def load_selected_ips(self):
        if os.path.exists("selected_ips.ini"):
            config = configparser.ConfigParser()
            config.read("selected_ips.ini")
            for section in config.sections():
                ip = config[section]["ip"]
                mac = config[section]["mac"]
                hostname = config[section]["hostname"]
                checkbox = QCheckBox()
                self.table_model.data.append((checkbox, ip, mac, hostname))
        self.table_model.layoutChanged.emit()
        self.update_labels()

    def load_settings(self):
        config = configparser.ConfigParser()
        key = load_key()
        if os.path.exists("settings.ini"):
            config.read("settings.ini")
            self.ssh_username = config.get("SSH", "username", fallback="")
            encrypted_password = config.get("SSH", "password", fallback="")
            if encrypted_password:
                self.ssh_password = decrypt_password(encrypted_password, key)
            else:
                self.ssh_password = ""

    def start_scan(self):
        self.table_model.data = []
        self.table_model.layoutChanged.emit()
        self.status_label.setText("状态: 正在扫描...")
        self.scan_thread.start()

    def display_results(self, results):
        self.table_model.data = [(QCheckBox(), client['ip'], client['mac'], client['hostname']) for client in results]
        self.table_model.layoutChanged.emit()
        self.update_labels()
        self.status_label.setText("状态: 扫描完成.")

    def update_labels(self):
        self.ip_count_label.setText(f"IP 数量: {self.table_model.rowCount()}")
        self.update_selected_count_label()

    def select_all(self):
        for row in range(self.table_model.rowCount()):
            index = self.table_model.index(row, 0)
            self.table_model.setData(index, Qt.Checked, Qt.CheckStateRole)
        self.update_labels()

    def deselect_all(self):
        for row in range(self.table_model.rowCount()):
            index = self.table_model.index(row, 0)
            self.table_model.setData(index, Qt.Unchecked, Qt.CheckStateRole)
        self.update_labels()

    def count_selected(self):
        count = 0
        for row in range(self.table_model.rowCount()):
            index = self.table_model.index(row, 0)
            if self.table_model.data[index.row()][0].checkState() == Qt.Checked:
                count += 1
        return count

    def save_data(self):
        selected_ips = []
        for row in range(self.table_model.rowCount()):
            index = self.table_model.index(row, 0)
            if self.table_model.data[index.row()][0].checkState() == Qt.Checked:
                selected_ips.append({
                    'ip': self.table_model.data[index.row()][1],
                    'mac': self.table_model.data[index.row()][2],
                    'hostname': self.table_model.data[index.row()][3],
                    'selected': True
                })

        if not selected_ips:
            QMessageBox.warning(self, "No Selection", "No rows selected to save.")
            return

        config = configparser.ConfigParser()
        for i, ip_info in enumerate(selected_ips, start=1):
            config[f'Host{i}'] = ip_info

        with open("selected_ips.ini", "w") as configfile:
            config.write(configfile)

        QMessageBox.information(self, "Save Successful", "Selected IPs saved to selected_ips.ini.")
        logging.info("Selected IPs saved successfully.")

    def update_selected_count_label(self):
        count = self.count_selected()
        self.selected_count_label.setText(f"选择的数量: {count}")

    def wake_on_lan(self):
        self.load_settings()  # 读取最新的设置信息
        for row in range(self.table_model.rowCount()):
            index = self.table_model.index(row, 0)
            if self.table_model.data[index.row()][0].checkState() == Qt.Checked:
                mac = self.table_model.data[index.row()][2]
                send_magic_packet(mac)
                logging.info(f"Sent WOL packet to MAC: {mac}")

    def shutdown(self):
        # 定义一个新的函数来执行实际的关机操作
       def shutdown_operation():
           self.load_settings()  # 读取最新的设置信息
           self.ssh_command("sudo poweroff")  #这个适合x86和800M
           self.ssh_command("sudo poweroff_wakeup")  #这个适合x86和800M
           sys.exit()  # 关闭程序

       # 创建一个新的线程来执行关机操作函数
       shutdown_thread = Thread(target=shutdown_operation)
       shutdown_thread.start()
        

    def restart(self):
        self.load_settings()  # 读取最新的设置信息
        self.ssh_command("sudo init 6")

    def ssh_command(self, command):
        for row in range(self.table_model.rowCount()):
            index = self.table_model.index(row, 0)
            if self.table_model.data[index.row()][0].checkState() == Qt.Checked:
                ip = self.table_model.data[index.row()][1]
                try:
                    conn = Connection(
                        host=ip,
                        user=self.ssh_username,
                        connect_kwargs={"password": self.ssh_password},
                    )
                    conn.run(command)
                    logging.info(f"Executed command '{command}' on {ip}")
                except Exception as e:
                    logging.error(f"SSH connection failed for {ip}: {e}")
                    print(f"SSH connection failed for {ip}: {e}")

    def open_settings(self):
        settings_dialog = SettingsDialog(self)
        if settings_dialog.exec_() == QDialog.Accepted:
            self.load_settings()  # 确保在设置对话框关闭后重新加载设置

if __name__ == "__main__":
    app = QApplication(sys.argv)
    window = MainWindow()
    window.show()
    sys.exit(app.exec_())
